import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier, ExtraTreesClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.svm import SVC
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import TimeSeriesSplit, cross_val_score
import mysql.connector

class FixedSSQPredictor:
    def __init__(self):
        self.mysql_config = {
            'host': 'localhost',
            'port': 3306,
            'user': 'root',
            'password': 'Duchenyi619',
            'database': 'redAndBlue'
        }
        
        # 初始化模型和标准化器
        self.models = {
            'red': {},
            'blue': {}
        }
        self.scalers = {
            'red_standard': StandardScaler(),
            'blue_standard': StandardScaler()
        }
    
    def load_data(self):
        """加载数据"""
        try:
            conn = mysql.connector.connect(**self.mysql_config)
            query = """
            SELECT draw_number, draw_date, 
                   red_ball_1, red_ball_2, red_ball_3, 
                   red_ball_4, red_ball_5, red_ball_6, 
                   blue_ball
            FROM lottery_data 
            ORDER BY draw_date ASC
            """
            df = pd.read_sql(query, conn)
            conn.close()
            
            # 转换红球为列表
            df['red_numbers'] = df[['red_ball_1', 'red_ball_2', 'red_ball_3', 
                                   'red_ball_4', 'red_ball_5', 'red_ball_6']].values.tolist()
            
            return df
        except Exception as e:
            print(f"数据加载失败: {e}")
            return pd.DataFrame()
    
    def preprocess_data(self, df):
        """数据预处理"""
        if len(df) < 50:
            print("数据量不足")
            return None, None, None, None
        
        # 特征提取
        features = []
        red_labels = []
        blue_labels = []
        
        sequence_length = 10
        
        for i in range(len(df) - sequence_length):
            # 提取特征
            sequence = df.iloc[i:i+sequence_length]
            next_draw = df.iloc[i+sequence_length]
            
            # 基础特征
            feature_vector = []
            
            # 红球特征
            for _, row in sequence.iterrows():
                red_nums = row['red_numbers']
                feature_vector.extend([
                    np.mean(red_nums) / 33,
                    np.std(red_nums) / 33,
                    sum(1 for x in red_nums if x % 2 == 1) / 6,  # 奇数比例
                    sum(1 for x in red_nums if x > 16) / 6,      # 大号比例
                ])
            
            # 蓝球特征
            for _, row in sequence.iterrows():
                blue_num = row['blue_ball']
                feature_vector.extend([
                    blue_num / 16,
                    blue_num % 2,  # 奇偶
                    blue_num > 8,  # 大小
                ])
            
            features.append(feature_vector)
            
            # 红球标签 (one-hot编码)
            red_label = np.zeros(33)
            for num in next_draw['red_numbers']:
                red_label[num-1] = 1
            red_labels.append(red_label)
            
            # 蓝球标签 (类别标签，不是one-hot)
            blue_labels.append(next_draw['blue_ball'] - 1)  # 转换为0-15
        
        return np.array(features), np.array(red_labels), np.array(blue_labels)
    
    def train_models(self):
        """训练模型"""
        print("开始训练模型...")
        
        # 加载数据
        df = self.load_data()
        if df.empty:
            print("数据加载失败")
            return False
        
        # 数据预处理
        features, red_labels, blue_labels = self.preprocess_data(df)
        if features is None:
            print("数据预处理失败")
            return False
        
        print(f"特征维度: {features.shape}")
        print(f"红球标签维度: {red_labels.shape}")
        print(f"蓝球标签维度: {blue_labels.shape}")
        
        # 时间序列交叉验证
        tscv = TimeSeriesSplit(n_splits=5)
        
        # 训练红球模型
        print("🔴 训练红球模型...")
        red_models = {
            'rf': RandomForestClassifier(n_estimators=100, max_depth=10, random_state=42),
            'gb': GradientBoostingClassifier(n_estimators=100, max_depth=6, random_state=42),
            'et': ExtraTreesClassifier(n_estimators=100, max_depth=10, random_state=42),
            'mlp': MLPClassifier(hidden_layer_sizes=(100, 50), random_state=42, max_iter=500),
            'svm': SVC(probability=True, random_state=42),
            'lr': LogisticRegression(random_state=42, max_iter=500)
        }
        
        self.models['red'] = red_models
        
        # 标准化特征
        features_scaled = self.scalers['red_standard'].fit_transform(features)
        
        # 训练每个红球位置
        for i in range(33):
            print(f"   训练红球 {i+1} 模型...")
            for name, model in red_models.items():
                model.fit(features_scaled, red_labels[:, i])
                
                # 交叉验证评估
                scores = cross_val_score(model, features_scaled, red_labels[:, i], 
                                       cv=tscv, scoring='accuracy')
                print(f"     {name} 交叉验证得分: {scores.mean():.4f} ± {scores.std():.4f}")
        
        # 训练蓝球模型
        print("🔵 训练蓝球模型...")
        blue_models = {
            'rf': RandomForestClassifier(n_estimators=100, max_depth=10, random_state=42),
            'gb': GradientBoostingClassifier(n_estimators=100, max_depth=6, random_state=42),
            'et': ExtraTreesClassifier(n_estimators=100, max_depth=10, random_state=42),
            'mlp': MLPClassifier(hidden_layer_sizes=(100, 50), random_state=42, max_iter=500),
            'svm': SVC(probability=True, random_state=42),
            'lr': LogisticRegression(random_state=42, max_iter=500)
        }
        
        self.models['blue'] = blue_models
        
        # 标准化特征
        features_scaled = self.scalers['blue_standard'].fit_transform(features)
        
        # 训练每个蓝球模型
        for name, model in blue_models.items():
            print(f"   训练蓝球 {name} 模型...")
            model.fit(features_scaled, blue_labels)
            
            # 交叉验证评估
            scores = cross_val_score(model, features_scaled, blue_labels, 
                                   cv=tscv, scoring='accuracy')
            print(f"     {name} 交叉验证得分: {scores.mean():.4f} ± {scores.std():.4f}")
        
        print("✅ 模型训练完成！")
        return True

if __name__ == "__main__":
    predictor = FixedSSQPredictor()
    success = predictor.train_models()
    if success:
        print("🎉 训练成功完成！")
    else:
        print("❌ 训练失败")
